{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0GTyDZmDVBml"
   },
   "source": [
    "# HyperStyle Inference Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1rvG7nRoGDUY"
   },
   "source": [
    "## Prepare Environment and Download HyperStyle Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QRzve9Y5DucV",
    "outputId": "86325ed1-7443-440a-bb24-4c3006dc3858"
   },
   "outputs": [],
   "source": [
    "#@title Clone HyperStyle Repo and Install Ninja { display-mode: \"form\" } \n",
    "import os\n",
    "os.chdir('/content')\n",
    "CODE_DIR = 'hyperstyle'\n",
    "\n",
    "## clone repo\n",
    "!git clone https://github.com/yuval-alaluf/hyperstyle.git $CODE_DIR\n",
    "\n",
    "## install ninja\n",
    "!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
    "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
    "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force\n",
    "\n",
    "os.chdir(f'./{CODE_DIR}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DxeMz-doE_pr"
   },
   "outputs": [],
   "source": [
    "#@title Import Packages { display-mode: \"form\" } \n",
    "import time\n",
    "import sys\n",
    "import pprint\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "sys.path.append(\".\")\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "from notebooks.notebook_utils import Downloader, HYPERSTYLE_PATHS, W_ENCODERS_PATHS, run_alignment\n",
    "from utils.common import tensor2im\n",
    "from utils.inference_utils import run_inversion\n",
    "from utils.domain_adaptation_utils import run_domain_adaptation\n",
    "from utils.model_utils import load_model, load_generator\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gKGnZ8DmICRU"
   },
   "source": [
    "## Define Download Configuration\n",
    "Select below whether you wish to download all models using `pydrive`. Note that if you do not use `pydrive`, you may encounter a \"quota exceeded\" error from Google Drive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QQ6XEmlHlXbk"
   },
   "outputs": [],
   "source": [
    "#@title { display-mode: \"form\" } \n",
    "download_with_pydrive = True #@param {type:\"boolean\"} \n",
    "downloader = Downloader(code_dir=CODE_DIR, use_pydrive=download_with_pydrive)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pU0snuo-SUHs"
   },
   "source": [
    "## Select Domain for Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XESWAO65kTJt"
   },
   "outputs": [],
   "source": [
    "#@title Select which domain you wish to perform inference on: { display-mode: \"form\" }\n",
    "experiment_type = 'faces' #@param ['faces', 'cars', 'afhq_wild']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9Tozsg81kTKA",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Define Inference Parameters\n",
    "\n",
    "Below we have a dictionary defining parameters such as the path to the pretrained model to use and the path to the image to perform inference on. While we provide default values to run this script, feel free to change as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2kE5y1-skTKC",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "EXPERIMENT_DATA_ARGS = {\n",
    "    \"faces\": {\n",
    "        \"model_path\": \"./pretrained_models/hyperstyle_ffhq.pt\",\n",
    "        \"w_encoder_path\": \"./pretrained_models/faces_w_encoder.pt\",\n",
    "        \"image_path\": \"./notebooks/images/face_image.jpg\",\n",
    "        \"transform\": transforms.Compose([\n",
    "            transforms.Resize((256, 256)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
    "    },\n",
    "    \"cars\": {\n",
    "        \"model_path\": \"./pretrained_models/hyperstyle_cars.pt\",\n",
    "        \"w_encoder_path\": \"./pretrained_models/cars_w_encoder.pt\",\n",
    "        \"image_path\": \"./notebooks/images/car_image.jpg\",\n",
    "        \"transform\": transforms.Compose([\n",
    "            transforms.Resize((192, 256)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
    "    },\n",
    "    \"afhq_wild\": {\n",
    "        \"model_path\": \"./pretrained_models/hyperstyle_afhq_wild.pt\",\n",
    "        \"w_encoder_path\": \"./pretrained_models/afhq_wild_w_encoder.pt\",\n",
    "        \"image_path\": \"./notebooks/images/afhq_wild_image.jpg\",\n",
    "        \"transform\": transforms.Compose([\n",
    "            transforms.Resize((256, 256)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])\n",
    "    }\n",
    "}\n",
    "\n",
    "EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pmOsCJWB6mGM",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Download Models\n",
    "To reduce the number of requests to fetch the model, we'll check if the model was previously downloaded and saved before downloading the model. We'll download the model for the selected experiment and save it to the folder `hyperstyle/pretrained_models`.\n",
    "\n",
    "We also need to verify that the model was downloaded correctly. All of our models should weigh approximately 1.3GB. Note that if the file weighs several KBs, you most likely encounter a \"quota exceeded\" error from Google Drive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jQ31J_m7kTJ8",
    "outputId": "97857f8a-45fe-40d4-c9b0-a7acd69bc8a0",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#@title Download HyperStyle Model { display-mode: \"form\" } \n",
    "if not os.path.exists(EXPERIMENT_ARGS['model_path']) or os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
    "    print(f'Downloading HyperStyle model for {experiment_type}...')\n",
    "    downloader.download_file(file_id=HYPERSTYLE_PATHS[experiment_type]['id'], file_name=HYPERSTYLE_PATHS[experiment_type]['name'])\n",
    "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
    "    if os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:\n",
    "        raise ValueError(\"Pretrained model was unable to be downloaded correctly!\")\n",
    "    else:\n",
    "        print('Done.')\n",
    "else:\n",
    "    print(f'HyperStyle model for {experiment_type} already exists!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bKeTTMU26mGN",
    "outputId": "19bf350a-4c17-4f2b-d20b-837b9349877d"
   },
   "outputs": [],
   "source": [
    "#@title Download WEncoder Model { display-mode: \"form\" } \n",
    "if not os.path.exists(EXPERIMENT_ARGS['w_encoder_path']) or os.path.getsize(EXPERIMENT_ARGS['w_encoder_path']) < 1000000:\n",
    "    print(f'Downloading the WEncoder model for {experiment_type}...')\n",
    "    downloader.download_file(file_id=W_ENCODERS_PATHS[experiment_type]['id'], file_name=W_ENCODERS_PATHS[experiment_type]['name'])\n",
    "    # if google drive receives too many requests, we'll reach the quota limit and be unable to download the model\n",
    "    if os.path.getsize(EXPERIMENT_ARGS['w_encoder_path']) < 1000000:\n",
    "        raise ValueError(\"Pretrained model was unable to be downloaded correctly!\")\n",
    "    else:\n",
    "        print('Done.')\n",
    "else:\n",
    "    print(f'WEncoder model for {experiment_type} already exists!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TAWrUehTkTKJ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load Pretrained Model\n",
    "We assume that you have downloaded all relevant models and placed them in the directory defined by the `EXPERIMENT_DATA_ARGS` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1t-AOhP1kTKJ",
    "outputId": "5d6cdb2a-0a8b-4516-a8f9-bed3be94bcb6",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#@title Load HyperStyle Model { display-mode: \"form\" } \n",
    "model_path = EXPERIMENT_ARGS['model_path']\n",
    "net, opts = load_model(model_path, update_opts={\"w_encoder_checkpoint_path\": EXPERIMENT_ARGS['w_encoder_path']})\n",
    "print('Model successfully loaded!')\n",
    "pprint.pprint(vars(opts))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4weLFoPbkTKZ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Define and Visualize Input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "r2H9zFLJkTKa",
    "outputId": "0bc3b3ea-388f-4161-9441-36c4883a8c8d",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "image_path = EXPERIMENT_DATA_ARGS[experiment_type][\"image_path\"]\n",
    "original_image = Image.open(image_path).convert(\"RGB\")\n",
    "if experiment_type == 'cars':\n",
    "    original_image = original_image.resize((192, 256))\n",
    "else:\n",
    "    original_image = original_image.resize((256, 256))\n",
    "original_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aTZcKMdK8y77",
    "outputId": "9b42f62c-a392-4d3e-f6e2-fb360ad7284a"
   },
   "outputs": [],
   "source": [
    "#@title Align Image (If Needed)\n",
    "input_is_aligned = False #@param {type:\"boolean\"}\n",
    "if experiment_type == \"faces\" and not input_is_aligned:\n",
    "    input_image = run_alignment(image_path)\n",
    "else:\n",
    "    input_image = original_image\n",
    "\n",
    "input_image.resize((256, 256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D0BmXzu1kTKg"
   },
   "source": [
    "## Perform Inference\n",
    "Now we'll run inference. By default, we'll run using 5 inference steps. You can change the parameter in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ct_jm0obOSDM",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#@title { display-mode: \"form\" } \n",
    "n_iters_per_batch = 5 #@param {type:\"integer\"}\n",
    "opts.n_iters_per_batch = n_iters_per_batch\n",
    "opts.resize_outputs = False  # generate outputs at full resolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ls5zb0fRkTKs",
    "outputId": "6b16d2b3-bf11-452b-ebdb-ef98ca0e7a02"
   },
   "outputs": [],
   "source": [
    "#@title Run Inference! { display-mode: \"form\" }\n",
    "img_transforms = EXPERIMENT_ARGS['transform']\n",
    "transformed_image = img_transforms(input_image) \n",
    "\n",
    "with torch.no_grad():\n",
    "    tic = time.time()\n",
    "    result_batch, result_latents, _ = run_inversion(transformed_image.unsqueeze(0).cuda(), \n",
    "                                                    net, \n",
    "                                                    opts,\n",
    "                                                    return_intermediate_results=True)\n",
    "    toc = time.time()\n",
    "    print('Inference took {:.4f} seconds.'.format(toc - tic))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nq0dkSz6kTKv"
   },
   "source": [
    "## Visualize Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VJpg06yzUU34"
   },
   "outputs": [],
   "source": [
    "def get_coupled_results(result_batch, transformed_image):\n",
    "    result_tensors = result_batch[0]  # there's one image in our batch\n",
    "    final_rec = tensor2im(result_tensors[-1]).resize(resize_amount)\n",
    "    input_im = tensor2im(transformed_image).resize(resize_amount)\n",
    "    res = np.concatenate([np.array(input_im), np.array(final_rec)], axis=1)\n",
    "    res = Image.fromarray(res)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 529
    },
    "id": "ca5BtxdUOSDN",
    "outputId": "50244361-6c4e-4cbb-e495-2fe818c64d3c"
   },
   "outputs": [],
   "source": [
    "if opts.dataset_type == \"cars\":\n",
    "    resize_amount = (256, 192) if opts.resize_outputs else (512, 384)\n",
    "else:\n",
    "    resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)\n",
    "\n",
    "res = get_coupled_results(result_batch, transformed_image)\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W5iecM0MUn1o"
   },
   "source": [
    "## Save Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5sY3hRjNUpiZ"
   },
   "outputs": [],
   "source": [
    "# save image \n",
    "outputs_path = \"./outputs\"\n",
    "os.makedirs(outputs_path, exist_ok=True)\n",
    "res.save(os.path.join(outputs_path, os.path.basename(image_path)))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Copy of inference_playground.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}